package org.xqh.test.weightbalance;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ThreadLocalRandom;


/**
 * @ClassName WeightedLoadBalancer
 * @Description TODO
 * @Author xuqianghui
 * @Date 2025/11/13 14:04
 * @Version 1.0
 */
public class WeightedLoadBalancer <T> {
    // 存储带权重的对象（泛型限制：对象需提供权重获取方式）
    private final List<T> objects;
    // 权重获取器（适配不同对象的权重字段）
    private final WeightExtractor<T> weightExtractor;
    // 权重总和（缓存，避免重复计算）
    private double totalWeight;
    // 轮询算法专用：当前偏移量（记录轮询状态）
    private double currentOffset = 0;

    /**
     * 权重提取器接口（适配不同对象的权重字段）
     */
    // 权重提取器改为返回double
    @FunctionalInterface
    public interface WeightExtractor<T> {
        double getWeight(T t);
    }

    /**
     * 构造器（初始化集合和权重提取逻辑）
     * @param objects 带权重的对象集合
     * @param weightExtractor 权重提取器（指定如何获取对象的权重）
     */
    public WeightedLoadBalancer(List<T> objects, WeightExtractor<T> weightExtractor) {
        // 校验集合非空
        if (objects == null || objects.isEmpty()) {
            throw new IllegalArgumentException("对象集合不能为空");
        }
        this.objects = new ArrayList<>(objects);
        this.weightExtractor = weightExtractor;
        // 计算权重总和并校验（至少有一个对象权重>0）
        this.totalWeight = calculateTotalWeight();
        if (totalWeight <= 0) {
            throw new IllegalArgumentException("至少有一个对象的权重需大于0");
        }
    }

    /**
     * 计算权重总和
     */
    private double calculateTotalWeight() {
        return objects.stream()
                .mapToDouble(weightExtractor::getWeight)
                .peek(weight -> {
                    if (weight < 0) {
                        throw new IllegalArgumentException("权重不能为负数");
                    }
                })
                .sum();
    }

    // -------------------------- 权重随机算法 --------------------------
    public T randomSelect() {
        // 生成 [0, totalWeight) 的随机数
        double randomNum = ThreadLocalRandom.current().nextDouble(totalWeight);
        double sum = 0;
        // 遍历累加权重，找到随机数所在的区间
        for (T obj : objects) {
            sum += weightExtractor.getWeight(obj);
            if (sum > randomNum) {
                return obj;
            }
        }
        // 理论上不会走到这里（除非权重计算异常）
        return objects.get(0);
    }

    // -------------------------- 权重轮询算法 --------------------------
    public T roundRobinSelect() {
        double maxWeight = 0;
        int selectedIndex = -1;

        // 1. 找到当前权重最大的对象（考虑偏移量）
        for (int i = 0; i < objects.size(); i++) {
            if(selectedIndex == -1){
                selectedIndex = i;
            }
            T obj = objects.get(i);
            double currentWeight = weightExtractor.getWeight(obj) + currentOffset;

            // 更新最大权重和选中索引
            if (currentWeight > maxWeight) {
                maxWeight = currentWeight;
                selectedIndex = i;
            }
        }

        // 2. 更新偏移量（选中对象的权重临时减去总权重，保证下一轮公平）
        T selectedObj = objects.get(selectedIndex);
        currentOffset = maxWeight - totalWeight;

        return selectedObj;
    }

    // -------------------------- 辅助方法 --------------------------
    /**
     * 刷新对象集合（如权重变更、对象增删后调用）
     */
    public void refreshObjects(List<T> newObjects) {
        this.objects.clear();
        this.objects.addAll(newObjects);
        this.totalWeight = calculateTotalWeight();
        this.currentOffset = 0; // 重置轮询状态
    }
}
